import json
import datasets
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from dotenv import load_dotenv
from logger import get_logger
import re
load_dotenv()

llm = ChatOpenAI(temperature=0.1, model="qwen2.5-72b-instruct")

logger = get_logger()

with open("MUSER-main/component_extra.json") as fin:
    component_list = json.load(fin)


dataset = datasets.load_from_disk("MUSER-main/muser-dataset")
# dataset = dataset.select(range(0, 160))


def make_prompt(fact, question):
    prompt = """\
下面给你一段放在三引号中的`法律文书事实描述`：
```
{fact}
```
请你根据上述案情描述，判断下面的问题是否成立：
`{question}`
你的回答应该按照以下步骤进行：
1.不要着急回答问题，先根据事实描述和已知信息进行分析;
2.根据分析得出你的结论;
3.最后再给出你的答案，以“【@是/否@】”的格式，格式化你的答案进行输出。
"""

    template = ChatPromptTemplate([
        ("system", "你是一个精通法律知识的助手，可以帮助人们分析法律文书，并回答相关问题。"),
        ("human", prompt),
    ])
    
    prompt_value = template.invoke(
        {
            "fact": fact,
            "question": question
        }
    )
    return prompt_value

def extract_from_resp(text):
    pattern = r"【@([^@]+)@】"
    match = re.search(pattern, text)

    if match:
        answer = match.group(1)  # 提取匹配的内容
        return answer
    else:
        return "否"

def dojudge(question, fact):
    prompt = make_prompt(fact, question)
    resp = llm.invoke(prompt).content
    # print('-' * 50)
    # print("问题：", question)
    # print("回答：", resp)
    res = extract_from_resp(resp)
    # print("抽取出的答案：", res)
    if "否" in res:
        return False
    else:
        return True

fout = open("feature-extra-72B.json", 'a+')

def process(data):
    res_dict = {}
    res_dict["idx"] = data["idx"]
    res_dict["index"] = data["index"]
    
    info_dict = {}
    fact = data["fact"]
    info_dict["fact"] = fact

    feature = [0] * len(component_list)
    
    for i, component in enumerate(component_list):
        if dojudge(component, fact):
            feature[i] = 1
    
    res_dict["feature"] = feature
    fout.write(json.dumps(res_dict, ensure_ascii=False) + "\n")
    fout.flush()

dataset.map(process)
fout.close()
